# 调用ocr 识别 original_img 中的图片， 识别出数据在在9*9 中的坐标
import cv2
import numpy as np
import torch
import torch.nn as nn

from main import handle

# 设置pytesseract的路径（如果需要）
# pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

print(torch.cuda.is_available())



def preprocess_image(image):
    # 转换为灰度图像
    image_ = np.copy(image)
    gray = cv2.cvtColor(image_, cv2.COLOR_BGR2GRAY)
    cv2.imshow("gray", gray)
    # 高斯模糊
    # blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    blurred = gray
    # 二值化处理
    _, thresh = cv2.threshold(blurred, 160, 255, cv2.THRESH_BINARY_INV) #  + cv2.THRESH_OTSU
    # 反转颜色
    thresh = 255 - thresh
    return thresh

def split_grid(image, rows=9, cols=9):
    height, width = image.shape
    cell_height = height // rows
    cell_width = width // cols
    # 将图像分割成单元格
    print(f"{height} {width} {cell_height} {cell_width}")
    cells = []
    for i in range(rows):
        for j in range(cols):
            # 计算高度范围
            y1 = i * cell_height
            y2 = (i + 1) * cell_height
            # 计算宽度范围
            x1 = j * cell_width
            x2 = (j + 1) * cell_width
            cell = image[y1:y2, x1:x2]

            cell[0:(cell_height//5),:] = 255
            cell[y2-(cell_height//5):, :] = 255
            cell[:, :(cell_width//5)] = 255
            cell[:, x2-(cell_width//5):] = 255

            cells.append(cell)
    return cells


def im_x_y_value(img,x, y, v):
    image_ = np.copy(img)
    height, width,_ = img.shape
    cell_height = height // 9
    cell_width = width // 9

    y1 = (x-1) * cell_height
    x1 = (y-1 )* cell_width
    y_ = y1 + cell_height//2
    x_ = x1 + cell_height//2
    #
    # 指定文本内容、起始坐标、字体、字体大小、字体颜色和字体厚度
    cv2.putText(image_, f'{v}', (x_, y_), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
    print((y_, x_))
    cv2.imshow("image", image_)
    return image_
 # 初始化ORB检测器
# orb = cv2.ORB_create()
def detect_digits(cells,image):
    digit_positions = []
    i__ = 0
    for i, cell in enumerate(cells):
        min_value = np.min(cell)
        row = i // 9
        col = i % 9
        if min_value< 150:
            i__ = i__+1
            # 调整图像大小以提高
            # 识别精度
            # cell = cv2.resize(cell, (28, 28))


            # print(f"识别结果: {digit_}, 置信度: {confidence:.2f}%")
            cv2.imshow("cell", cell)
            key = cv2.waitKey(0)
            print(f"{key}   {chr(key)}")
            image = im_x_y_value(image,row+1,col+1,int(chr(key)))
            digit_positions.append((row, col, int(chr(key))))
        else:
            digit_positions.append((row, col, int(-1)))

    print(f"{i__ = }")
    return digit_positions

# 读取图像
image = cv2.imread(r"./original_img/002.png")
# 预处理图像
preprocessed = preprocess_image(image)

cv2.imshow("image", image)
cv2.imshow("preprocessed", preprocessed)
# cv2.waitKey(0)
# 分割网格
cells = split_grid(preprocessed)

print(len(cells))

# 检测数字
digit_positions = detect_digits(cells,image)

# 输出数字所在位置
value_list=[]
for row, col, digit in digit_positions:
    if digit != -1:
        print(f"数字 {digit} 位于第 {row + 1} 行，第 {col + 1} 列")
        index = (row + 1) *100 + (col + 1)*10 + digit
        value_list.append(index)

result = handle(value_list)
# cv2.imshow("image2222", image)
# cv2.waitKey(0)
for i in range(9):
    for j in range(9):
        v = result[i][j]
        image = im_x_y_value(image, i + 1, j + 1, int(v))
cv2.imshow("image", image)
cv2.waitKey(0)
print(result)


